import torch
import torch.nn as nn
import torch.nn.functional as F


class CustomLoss(nn.Module):
    """
    自定义损失函数, 这里使用MSELoss
    """

    def __init__(self):
        super(CustomLoss, self).__init__()

    def forward(self, preds, targets):
        preds = preds.float()
        targets = targets.float()
        loss = F.mse_loss(preds, targets, reduction="none")
        return loss.mean()


# 如果您要测试这个损失函数：
if __name__ == "__main__":
    # 随机生成模拟数据
    preds = torch.rand(4, 5, 64, 64)  # 假设有4张图片，5个关键点，每个热力图大小为64x64
    targets = torch.rand(4, 5, 64, 64)

    criterion = CustomLoss()
    loss = criterion(preds, targets)
    print(loss)
